package com.guchenbo.springboot;

import java.io.IOException;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 * @author guchenbo
 * @date 2022/6/27
 */
@Component
public class DemoWebSocketHandler extends TextWebSocketHandler {
    protected static final Logger log = LoggerFactory.getLogger(DemoWebSocketHandler.class);

    private final Map<String, Set<String>> userSession = new ConcurrentHashMap<>();
    private final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    @Resource
    private DemoService demoService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        super.afterConnectionEstablished(session);
        String user = getUser(session);
        String id = session.getId();
        log.info("conn established for {}-{}", user, id);

        // 可能多个连接使用一个用户
        Set<String> ids = getUserIds(user);
        ids.add(id);
        userSession.put(user, ids);

        sessionMap.putIfAbsent(id, session);
        sendMsg(user, demoService.count());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        super.afterConnectionClosed(session, status);
        String user = getUser(session);
        String id = session.getId();

        Set<String> ids = getUserIds(user);
        ids.remove(id);

        sessionMap.remove(id);
        session.close();
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        super.handleTextMessage(session, message);
        String user = getUser(session);
        log.info("receive msg {}-{}", user, message.getPayload());
    }

    public void sendMsg(String user, String count) throws IOException {
        Set<String> ids = getUserIds(user);
        TextMessage message = new TextMessage(count);

        ids.forEach(id -> {
            WebSocketSession session = sessionMap.get(id);
            if (session != null) {
                try {
                    session.sendMessage(message);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        });
    }

    private String getUser(WebSocketSession session) {
        Map<String, Object> att = session.getAttributes();
        return (String) att.get("user");
    }

    private Set<String> getUserIds(String user) {
        return userSession.getOrDefault(user, new HashSet<>());
    }
}
